import json

from zhipuai import ZhipuAI


class ChatGLMPRO:
    def __init__(self, api_key="df13780c7060b03a47628f3b71cf24b1.BqJqARhQJDgRdhwj"):
        self.client = ZhipuAI(api_key=api_key)

    def my_gradio_response(self, ask, history, stream=True):
        system_template = ask
        chatmsg = []
        for chat in history:
            chatmsg.append({
                "role": "user",
                "content": str(chat[0])
            })
            chatmsg.append({
                "role": "assistant",
                "content": str(chat[1])
            })

        chatmsg.append({
            "role": "user",
            "content": str(system_template)
        })

        response = self.client.chat.completions.create(
            model="glm-4",
            messages=chatmsg,
            temperature=0.95,
            top_p=0.7,
            stream=stream
        )

        return response

    def answer(self, ask, org_ask, history):
        # org_ask代表要显示的ask
        response = self.my_gradio_response(ask, history, stream=True)
        history.append([org_ask, ""])
        for chunk in response:
            history[-1][1] += chunk.choices[0].delta.content
            yield history

    def answer_no_stream(self, ask, history):
        response = self.my_gradio_response(ask, history, stream=False)
        return response.choices[0].message.content
